# 根据长度要求智能选择示例
# 根据输入相似度选择示例（最大边际相关性）
# 根据输入相似度选择示例（最大余弦相似度）
from langchain.prompts import PromptTemplate
from langchain.prompts import FewShotPromptTemplate
# 根据输入的提示词长度综合计算最终长度，智能截取或者添加提示词的示例
from langchain.prompts import LengthBasedExampleSelector
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.example_selectors import MaxMarginalRelevanceExampleSelector
from langchain_ollama import OllamaEmbeddings

# 假设已经有这么多的提示词
examples = [
    {"input": "happy", "output": "sad"},
    {"input": "tall", "output": "short"},
    {"input": "sunny", "output": "gloomy"},
    {"input": "windy", "output": "calm"},
    {"input": "高兴", "output": "悲伤"}
]

# 构造提示词模版
example_prompt = PromptTemplate(
    input_variables=["input", "output"],
    template="原词: {input} \n反义词：{output}",
)

# 调用长度示例选择器
example_selector = LengthBasedExampleSelector(
    # 传入提示词示例组
    examples=examples,
    # 传入提示词模版
    example_prompt=example_prompt,
    # 设置格式化后提示词的最大长度
    max_length=25
    # 内置的get_text_length,如果默认分词计算方式不满足，可以扩展
    # 默认计算方式：get_text_length: Callable[[str],int] = lambda x:len(re.split("\n"))
)

# 使用小样本提示词模版来实现动态示例的调用
dynamic_prompt = FewShotPromptTemplate(
    example_selector=example_selector, # 注意
    example_prompt=example_prompt,
    prefix="给出每个输入词的反义词",
    suffix="原词: {adjective}\n反义词：",
    input_variables=["adjective"]
)
print(dynamic_prompt.format(adjective="big"))


# MaxMarginalRelevanceExampleSelector根据与输入最相似的示例的组合来选择示例，同时优化多样性。它通过找到嵌入与输入具有最大余弦相似性的示例来做到这一点，然后迭代地添加它们，同时惩罚它们与已经选择的示例的接近程度。
"""
LengthBasedExampleSelector使用示例列表和嵌入创建k-shot示例选择器，以下为几个变量：
- examples: List[dict] # 提示模板所需的示例列表。
- embeddings: Embeddings # 初始化的嵌入API接口，例如OpenAIEmbeddings()。
- vectorstore_cls (Type[VectorStore]) # A vector store DB interface class, e.g. FAISS.
- k: int=4 # 选择生成的示例个数
"""

embeddings = OllamaEmbeddings(
    model="shaw/dmeta-embedding-zh",
    base_url="http://10.2.4.31:11434",
)
example_selector = MaxMarginalRelevanceExampleSelector.from_examples(
    # The list of examples available to select from.
    examples,
    embeddings=embeddings,
    # 存储embedding
    vectorstore_cls=FAISS,
    # The number of examples to produce.
    k=2,
)

mmr_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="Give the antonym of every input",
    suffix="Input: {adjective}\nOutput:",
    input_variables=["adjective"],
)

# 输入是一种感觉，所以应该选择快乐/悲伤的例子作为第一个
print(mmr_prompt.format(adjective="worried"))

